#!/usr/bin/python
#Programmer: Aaron Steele
#Purpose: Convert the codon frequences and calculate the min-max of a pfam family

import sys
import os
import string
import pfam_mm as mm
import subprocess

if __name__ == "__main__":

	start_dir = sys.argv[1]
	for dirname, dirnames, filenames in os.walk(start_dir):
	
		#For each pfam-direcotry find the fasta and cf file
		for subdirname in dirnames:
			fam_name = subdirname.replace("/","")
			print ("Processing family: " + fam_name)

			fa_file = start_dir + "/" + subdirname + "/" + fam_name + ".fa"
			cf_file = start_dir + "/" + subdirname + "/" + fam_name + ".cf"

			mm.single_mm(fa_file, cf_file)

			#Move the file
			command = "mv tmp.mm " + start_dir + "/" + subdirname + "/" + fam_name + ".mm"
			proc = subprocess.Popen(command,shell=True)		 
			proc.wait()

	#Clean it up
	command = "rm mm.tmp.fasta"
	proc = subprocess.Popen(command,shell=True)
	proc.wait()
	command = "rm tmp.cf"
	proc = subprocess.Popen(command,shell=True)
	proc.wait()
	
def single_mm(fa_file, cf_file):
	fasta_in = open(fa_file,"r")
	cf_in = open(cf_file,"r")	

	ids = []
	cfs = []

	#Read in all the data from the codon frequencey file
	line = cf_in.readline()
	while line:
		line = line.rstrip()

		if line.startswith(">"):
			line = line.replace(">","")
			ssplit = line.split(":")
			ids.append(ssplit[1])
		else:
			cfs.append(line)

		line = cf_in.readline()
	cf_in.close()


	#Read through the sequence data
	line = fasta_in.readline()
	while line:
		line = line.rstrip()

		#For each sequence extract the sequence
		if line.startswith(">"):
			line = line.replace(">","")
			full_id = line
			ssplit = line.split(":")
			id = ssplit[1].rstrip()
			line = fasta_in.readline()
			seq = line.rstrip()
			
			if id in ids:	
				#Export necessary files
				mm.export_seq(full_id,seq)
				index = ids.index(id)
				mm.convert_export_codons(cfs[index])
		
				#Start a subprocess for each id
				command = "perl rc-min-max.pl mm.tmp.fasta tmp.cf >> tmp.mm"
				proc = subprocess.Popen(command,shell=True)
				proc.wait()
			#else:
				#print("Bad ID: " + id)
		line = fasta_in.readline()

	fasta_in.close()
	cf_in.close()

def export_seq(id, seq):
	mm_tmp = open("mm.tmp.fasta","w")
	mm_tmp.write(">" + id + "\n")
	mm_tmp.write(seq + "\n")
	mm_tmp.close()


def convert_export_codons(codon_freq):
	#Codon order from Kazusa
	codon_strings = ["CGA", "CGC", "CGG", "CGT", "AGA", "AGG", "CTA", "CTC", "CTG", "CTT", "TTA", "TTG", "TCA", "TCC", "TCG", "TCT", "AGC", "AGT", "ACA", "ACC", "ACG", "ACT", "CCA", "CCC", "CCG", "CCT", "GCA", "GCC", "GCG", "GCT", "GGA", "GGC", "GGG", "GGT", "GTA", "GTC", "GTG", "GTT", "AAA", "AAG", "AAC", "AAT", "CAA", "CAG", "CAC", "CAT", "GAA", "GAG", "GAC", "GAT", "TAC", "TAT", "TGC", "TGT", "TTC", "TTT", "ATA", "ATC", "ATT", "ATG", "TGG", "TAA", "TAG", "TGA"]

	ssplit = codon_freq.split()
	cf_tmp = open("tmp.cf","w")

	#Get the sum of all codons used
	sum = 0
	for i in range(len(ssplit)):
		sum += int(ssplit[i])

	#Output the codon frequency in the messed up format that is required	
	count = 1
	for j in range(len(ssplit)):
		freq = float(ssplit[j]) / sum * 1000
		if count == 4:
			cf_tmp.write(codon_strings[j] + " " + str(freq) + "(" + ssplit[j] + ")" + "\n")
			count = 1
		else:
			cf_tmp.write(codon_strings[j] + " " + str(freq) + "(" + ssplit[j] + ")" + "\t")
			count +=1

	cf_tmp.close()
		
	
